import os.path
import time
from multiprocessing.resource_tracker import register

import NetEngine
import NetProtocol.out.python.net_protocol as net_protocol


class Client:
    def __init__(self, engine: NetEngine.NetEngineServer, session_p: NetEngine.SessionHandle):
        self._engine = engine
        self._session_p = session_p
        self._is_login = False
        self._protocol = net_protocol.NetProtocol()
        self._rx_buf = []
        pass

    def _allocate_uid(self) -> int:
        for i in range(0xffffffff):
            if not os.path.exists('~\\.tchatserver\\%d.pem' % i):
                return i
        return -1

    def _save_cert(self, uid, new_cert: bytes):
        with open('~\\.tchatserver\\%d.pem' % uid, mode='wb') as fp:
            fp.write(new_cert)

    def recv_data(self, data: bytes):
        self._rx_buf += list(data)
        frame = self._protocol.cut_frame(self._rx_buf)
        if frame is not None:
            # 先看看消息类型
            msg_head = self._protocol.decode_msg_head(frame)
            if self._is_login:
                pass
            else:  # 未登入时，只处理注册公钥和登入消息
                if msg_head.MsgType() == net_protocol.MsgTypeRegisterPublicKey:  # 如果是注册公钥，存注册的公钥，还得分配一个uid
                    msg = self._protocol.decode_msg(frame)
                    new_uid = self._allocate_uid()
                    if new_uid >= 0:
                        self._save_cert(new_uid, bytes(msg.PublicKey))
                        register_result = 0
                    else:
                        print('UID分配失败')
                        register_result = 1
                    msg_register_pubkey_result = net_protocol.MsgRegisterPublicKeyResult(new_uid)
                    msg_register_pubkey_result.RegisterResult = register_result
                    msg_register_pubkey_result.NewUID = new_uid
                    # 将服务器公钥返回给客户端
                    with open('~\\.tchatserver\\pub_key.pem', mode='rb') as fp:
                        server_pub_key = fp.read()
                    msg_register_pubkey_result.ServerPublicKey = list(server_pub_key)
                    msg_register_pubkey_result.ServerPublicKeyLen = len(msg_register_pubkey_result.ServerPublicKey)
                    ret_frame = self._protocol.encode_msg(msg_register_pubkey_result)
                    self._engine.send(self._session_p, ret_frame)
                elif msg_head.MsgType() == net_protocol.MsgTypeLogin:  # 如果是登入，需要检查对应的uid证书是否存在，并要初始化协议栈加解密功能
                    msg_login_result = net_protocol.MsgLoginResult(msg_head.UID())
                    if os.path.exists('~\\.tchatserver\\%d.pem' % msg_head.UID()):
                        self._protocol.init_asymmetric('~\\.tchatserver\\%d.pem' % msg_head.UID(), '~\\.tchatserver\\pri_key.pem')
                        msg = self._protocol.decode_msg(frame)
                        if msg is not None:
                            self._protocol.init_symmetric(bytes(msg.SymmetricKey))
                            msg_login_result.LoginResult = 0  # 成功
                            msg_login_result.SetEncryptMode(net_protocol.EncModeAsymmetric)
                            self._is_login = True
                        else:
                            msg_login_result.LoginResult = 1  # 解密失败
                            msg_login_result.SetEncryptMode(net_protocol.EncModeAsymmetric)
                            print('解密登入消息失败')
                    else:
                        msg_login_result.LoginResult = 2  # 没有证书
                        msg_login_result.SetEncryptMode(net_protocol.EncModeNone)
                        print('没有对应客户端证书')
                    ret_frame = self._protocol.encode_msg(msg_login_result)
                    self._engine.send(self._session_p, ret_frame)
        pass

    def run_once(self):
        pass


class TChatServer:
    def __init__(self):
        self.client_dir: dict[NetEngine.SessionHandle, Client] = {}
        self.engine = NetEngine.NetEngineServer()
        self.engine.register_data_notify(self.data_notify)
        self.engine.register_server_event_notify(self.event_notify)
        if not os.path.exists('~\\.tchatserver\\pri_key.pem'):
            if not os.path.exists('~\\.tchatserver'):
                os.makedirs('~\\.tchatserver')
            protocol = net_protocol.NetProtocol()
            protocol.enc_module.generate_asymmetric_key()
            pri_key = protocol.enc_module.export_asymmetric_local_private_key()
            with open('~\\.tchatserver\\pri_key.pem', mode='wb') as fp:
                fp.write(pri_key)
            pub_key = protocol.enc_module.export_asymmetric_local_public_key()
            with open('~\\.tchatserver\\pub_key.pem', mode='wb') as fp:
                fp.write(pub_key)
        self.engine.start_listen(50200)

    def data_notify(self, session_p: NetEngine.SessionHandle, is_recv: bool, data: bytes):
        client = self.client_dir[session_p]
        if client is not None:
            client.recv_data(data)
        pass

    def event_notify(self, session_p: NetEngine.SessionHandle, event: NetEngine.ServerEvent):
        print('server event', session_p, event)
        if event == NetEngine.ServerEvent.NewSession:
            print('new session', session_p)
            self.client_dir[session_p] = Client(self.engine, session_p)
        elif event == NetEngine.ServerEvent.SessionDone:
            print('del session', session_p)
            del self.client_dir[session_p]

    def run(self):
        while True:
            for client in self.client_dir.values():
                client.run_once()
            time.sleep(1)


server = TChatServer()
server.run()
